from pyNastran.op2.op2 import OP2
import json


# README
#
#
#
#
#

class Op2Transformer:
    model: None
    displacements: None

    def __init__(self, path):
        # 具体的结构查看readme
        self.displacements = {}
        self.stresses = {}

        self.path = path
        op2 = OP2()
        op2.read_op2(path)
        self.model = op2
        self.get_status()
        self.get_displacements()
        self.get_stress()
        self.get_strain()

    def get_status(self):
        # 获取op2文件的统计信息
        op2_stats = self.model.get_op2_stats(False)
        # 打印op2文件的统计信息
        print("OP2 File Statistics:\n")
        print(op2_stats.split("\n"))

    def get_displacements(self):
        # 获取位移结果
        displacements = self.model.displacements

        # 遍历每个子案例的位移结果
        for subcase, disp in displacements.items():
            print(f"Subcase: {subcase}")
            self.displacements[int(subcase)] = {
                "node_ids": [],
                "t": [],
                "r": [],
            }

            # 获取节点编号和位移值
            node_ids = disp.node_gridtype[:, 0].tolist()
            tx = disp.data[0, :, 0].tolist()
            ty = disp.data[0, :, 1].tolist()
            tz = disp.data[0, :, 2].tolist()
            rx = disp.data[0, :, 3].tolist()
            ry = disp.data[0, :, 4].tolist()
            rz = disp.data[0, :, 5].tolist()

            # 打印节点编号和位移值
            for i, node_id in enumerate(node_ids):
                self.displacements[subcase]["node_ids"].append(node_id)
                self.displacements[subcase]["t"].append(tx[i])
                self.displacements[subcase]["t"].append(ty[i])
                self.displacements[subcase]["t"].append(tz[i])
                self.displacements[subcase]["r"].append(rx[i])
                self.displacements[subcase]["r"].append(ry[i])
                self.displacements[subcase]["r"].append(rz[i])

    def get_stress(self):
        # 获取应力结果
        if hasattr(self.model, "ctria3_stress"):
            self.parse_ctria3_stress()
        if hasattr(self.model, "cquad4_stress"):
            self.parse_cquad4_stress()

    def parse_ctria3_stress(self):
        ctria3_stress = self.model.ctria3_stress
        # print("ctria3_stress", ctria3_stress)
        # 遍历每个子案例的应力结果
        for subcase, stress_data in ctria3_stress.items():
            if int(subcase) not in self.stresses:
                self.stresses[int(subcase)] = {}
            # 🙂  CTRIA3 此处为了和标准写法一样
            self.stresses[int(subcase)]["CTRIA3"] = {
                # 获取单元编号和节点编号
                "element_ids": stress_data.element_node[:, 0].tolist(),
                "node_ids": stress_data.element_node[:, 1].tolist(),
                # 获取应力值
                "fiber_distance": stress_data.data[0, :, 0].tolist(),
                "oxx": stress_data.data[0, :, 1].tolist(),
                "oyy": stress_data.data[0, :, 2].tolist(),
                "txy": stress_data.data[0, :, 3].tolist(),
                "angle": stress_data.data[0, :, 4].tolist(),
                "omax": stress_data.data[0, :, 5].tolist(),
                "omin": stress_data.data[0, :, 6].tolist(),
                "von_mises": stress_data.data[0, :, 7].tolist(),
            }

    def parse_cquad4_stress(self):
        cquad4_stress = self.model.cquad4_stress
        # print("cquad4_stress", cquad4_stress)
        # 遍历每个子案例的应力结果
        for subcase, stress_data in cquad4_stress.items():
            if int(subcase) not in self.stresses:
                self.stresses[int(subcase)] = {}
            self.stresses[int(subcase)]["CQUAD4"] = {
                # 获取单元编号和节点编号
                "element_ids": stress_data.element_node[:, 0].tolist(),
                "node_ids": stress_data.element_node[:, 1].tolist(),
                # 获取应力值
                "fiber_distance": stress_data.data[0, :, 0].tolist(),
                "oxx": stress_data.data[0, :, 1].tolist(),
                "oyy": stress_data.data[0, :, 2].tolist(),
                "txy": stress_data.data[0, :, 3].tolist(),
                "angle": stress_data.data[0, :, 4].tolist(),
                "omax": stress_data.data[0, :, 5].tolist(),
                "omin": stress_data.data[0, :, 6].tolist(),
                "von_mises": stress_data.data[0, :, 7].tolist(),
            }

    def get_strain(self):
        print("todo")

    def to_json(self) -> str:
        data = {
            "displacements": self.displacements,
            "stresses": self.stresses,
        }
        return json.dumps(data)
